"""测试一元线性回归"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from linear_regression import LinearRegression

"""获取数据"""
data = pd.read_csv("../data/world-happiness-report-2017.csv")
train_data = data.sample(frac=0.8)  # 设置训练集和测试集的比例是0.8
test_data = data.drop(train_data.index)  # 去掉训练集的索引

"""输入特征名字"""
input_param_name = "Economy..GDP.per.Capita."
output_param_name = "Happiness.Score"

x_train = train_data[[input_param_name]].values
y_train = train_data[[output_param_name]].values


x_test = test_data[[input_param_name]].values
y_test = test_data[[output_param_name]].values

"""做个图看一眼"""
plt.scatter(x_train, y_train, label="Train data")
plt.scatter(x_test, y_test, label="Test data")
plt.xlabel(input_param_name)
plt.ylabel(output_param_name)
plt.title("Happy")
plt.legend()
plt.show()

"""开始线性回归"""
num_iteration = 500  # 迭代次数
learn_rate = 0.01  # 学习率

linear = LinearRegression(x_train, y_train)  # 实例化
(theta, cost_history) = linear.train(learn_rate, num_iteration)
print("开始时的损失", cost_history[0])
print("训练后的损失", cost_history[-1])

"""画出来损失函数的趋势"""
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
plt.plot(range(len(cost_history)), cost_history)
plt.xlabel("迭代次数")
plt.ylabel("损失")
plt.title("梯度下降")
plt.show()

"""测试回归方程"""
predict_num = 100
x_predictions = np.linspace(x_train.min(), x_train.max(), predict_num).reshape(predict_num, 1)
y_predictions = linear.prodict(x_predictions)

"""画出测试集的点，和回归曲线"""
plt.scatter(x_train, y_train, label="Train data")
plt.scatter(x_test, y_test, label="Test data")
plt.plot(x_predictions, y_predictions, 'r', label="回归方程")
plt.xlabel(input_param_name)
plt.ylabel(output_param_name)
plt.title("Happy")
plt.legend()
plt.show()

